from diffusers import DiffusionPipeline
import torch
import os
dataRoot = os.path.join(os.path.dirname(os.path.abspath(__file__)),'../DownloadRoot')

pipe = DiffusionPipeline.from_pretrained(os.path.join(dataRoot,"models/stabilityai/stable-diffusion-xl-base-1.0"), torch_dtype=torch.float16, variant="fp16")

pipe.to("cuda")

prompt = "Watercolor painting of a desert landscape, with sand dunes, mountains, and a blazing sun, soft and delicate brushstrokes, warm and vibrant colors"

negative_prompt = "(EasyNegative),(watermark), (signature), (sketch by bad-artist), (signature), (worst quality), (low quality), (bad anatomy), NSFW, nude, (normal quality)"

seed = torch.Generator("cuda").manual_seed(42)

# 首先运行Base模型的Pipeline，输出格式为output_type="latent"
image = pipe(prompt=prompt, negative_prompt=negative_prompt, generator=seed, output_type="latent").images

# 构建Stable Diffusion XL Refiner模型的Pipeline，加载Stable Diffusion XL Refiner模型
pipe = DiffusionPipeline.from_pretrained(os.path.join(dataRoot,"models/stabilityai/stable-diffusion-xl-refiner-1.0"), torch_dtype=torch.float16, variant="fp16")
# "本地路径/stable-diffusion-xl-refiner-0.9"表示我们需要加载的Stable Diffusion XL Refiner模型，
# 大家可以关注Rocky的公众号WeThinkIn，后台回复：SDXL模型，即可获得资源链接

pipe.to("cuda")

images = pipe(prompt=prompt, negative_prompt=negative_prompt, generator=seed, image=image).images

images[0].save("test.png")